import env
import gym
import random
from transformers import AutoModel, AutoTokenizer, AutoConfig, T5ForConditionalGeneration
from itertools import combinations

actions = ['{},{}'.format(i, j) for i, j in combinations(list(range(20)), 2)]

def main():
    tokenizer = AutoTokenizer.from_pretrained('/home/boai/LLM_models/T5_small', trust_remote_code=True)
    model = T5ForConditionalGeneration.from_pretrained('/home/boai/LLM_models/T5_small', trust_remote_code=True)
    env = gym.make('TSP_reverse-v0')
    for _ in range(5):
        env.reset()
        prompt = env.get_description()
        # prompt = 'How is the weather?'
        input_ids = tokenizer(prompt, return_tensors='pt').input_ids
        # print(input_ids)
        action = random.choice(actions)
        # action = 'How is the weather?'
        decoder_input_ids = tokenizer(action, return_tensors='pt').input_ids
        outputs = model.generate(input_ids=input_ids, decoder_input_ids=decoder_input_ids)
        print(f"input:{prompt}\noutput:{tokenizer.decode(outputs[0], skip_special_tokens=True)}\n'")

if __name__ == '__main__':
    main()